#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import opengate as gate
import linac_elekta_synergy
from pathlib import Path

if __name__ == "__main__":
    # create the simulation
    sim = gate.Simulation()

    # main options
    sim.g4_verbose = False
    # sim.visu = True
    sim.visu_type = "vrml"
    sim.check_volumes_overlap = False
    sim.number_of_threads = 1
    sim.random_seed = "auto"
    sim.progress_bar = True
    sim.output_dir = "./output"

    # units
    m = gate.g4_units.m
    cm = gate.g4_units.cm
    mm = gate.g4_units.mm
    nm = gate.g4_units.nm
    Bq = gate.g4_units.Bq
    MeV = gate.g4_units.MeV

    # world size
    sim.world.size = [1 * m, 1 * m, 3 * m]
    sim.world.material = "G4_AIR"

    # add a linac
    linac = linac_elekta_synergy.add_linac(sim, "my_linac")
    linac.translation = [0, 0, 100 * cm]

    # virtual plane for phase space
    plane = sim.add_volume("Tubs", "phase_space_plane")
    plane.mother = linac.name
    plane.material = "G4_AIR"
    plane.rmin = 0
    plane.rmax = 100 * mm
    plane.dz = 1 * nm  # half height
    plane.translation = [0, 0, -27 * cm]
    plane.color = [1, 0, 0, 1]  # red

    # e- source
    source = sim.add_source("GenericSource", "electron_source")
    source.particle = "e-"
    source.attached_to = f"{linac.name}_target"
    source.energy.type = "gauss"
    source.energy.mono = 6.7 * MeV
    source.energy.sigma_gauss = 0.077 * MeV
    source.position.type = "disc"
    source.position.sigma_x = 1.274 * mm
    source.position.sigma_y = 1.274 * mm
    source.position.translation = [0, 0, 0.6 * mm]
    source.direction.type = "momentum"
    source.direction.momentum = [0, 0, -1]
    # source.n = 1000000 / sim.number_of_threads
    source.n = 10000 / sim.number_of_threads

    # add stat actor
    stats = sim.add_actor("SimulationStatisticsActor", "stats")
    stats.track_types_flag = True
    stats.output_filename = "stats2.txt"

    # PHSP plane
    phsp = sim.add_actor("PhaseSpaceActor", "PhaseSpace")
    phsp.attached_to = plane.name
    phsp.attributes = [
        "KineticEnergy",
        "Weight",
        "PrePositionLocal",
        "PreDirectionLocal",
    ]
    phsp.output_filename = "phsp.root"
    phsp.debug = False
    f = sim.add_filter("ParticleFilter", "f")
    f.particle = "gamma"
    phsp.filters.append(f)

    # phys
    sim.physics_manager.physics_list_name = "G4EmStandardPhysics_option3"
    sim.physics_manager.set_production_cut("world", "all", 1 * mm)

    # splitting in the target : need to create a region first
    region_linac_target = sim.physics_manager.add_region(name=f"{linac.name}_target")
    region_linac_target.associate_volume(linac)
    # split the e- Bremstralhung 100X times (if the energy is lower than 10 MeV)
    s = f"/process/em/setSecBiasing eBrem {linac.name}_target 100 10 MeV"
    sim.g4_commands_after_init.append(s)

    # ---------------------------------------------------------------------
    # start simulation
    sim.run()

    # end
    print(stats)
